from openopt import NLP
from numpy import cos, arange, ones, asarray, abs, zeros, sqrt, asscalar, array
from pylab import legend, show, plot, subplot, xlabel, subplots_adjust
from string import rjust, ljust, expandtabs

import JT
import Calibration

#===============================================================================
#
#    Data Paths
#
#===============================================================================

fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
movingImageName = 'C:/Users/bryan/bryan-code/trunk/Images/CalibratedDRRImage.mhd'
inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
calFileLat = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'


#===============================================================================
#
#    Setup DRR
#
#===============================================================================

xraycam = JT.XRayCamera()
drr = JT.DRR()
drr.SetXRayCamera(xraycam)
drr.SetVolumeFilename(inputVolFileName)
drr.SetBlackOnWhiteImage()
drr.InteractionOff()

cal = Calibration.ExternalCalibration()
cal.LoadConsolidatedCalibration(calFileLat)
cal.LoadStaFile(staFile)

xraycam.SetExternalCalibration(cal)

vtkMatrix = cal.ConvertNumpyToVTKArray(cal._VolumeT)
drr.SetVolumeUserMatrix(vtkMatrix)
drr.SetBlackOnWhiteImage()

# Set Fixed image using path
fixedImage = JT.FixedImage()
fixedImage.SetFileName(fixedImageName)

reg=JT.NewRegistration()
reg.SetFixedImage(fixedImage.GetImage(0))
reg.SetMovingImageGenerator(drr)

f=reg.GetMetric

N = 6
M = 5

x0=array((0,0,0,2.5,-5,5))

lb = -5*ones(N)
ub = 5*ones(N)

colors = ['b', 'k', 'y', 'r', 'g']
#############
#solvers = ['ralg','scipy_cobyla', 'lincher']
solvers = ['ralg', 'scipy_cobyla', 'lincher','ipopt','algencan' ]
solvers = ['ralg', 'ralg3', 'ralg5']
solvers = ['ralg',  'scipy_cobyla']
#solvers = ['ipopt']
#############
colors = colors[:len(solvers)]

lines, results = [], {}
for j in range(len(solvers)):
    solver = solvers[j]
    color = colors[j]
    p = NLP(f, x0, lb = lb, ub = ub, ftol = 1e-6, maxFunEvals = 1e7, maxIter = 120, plot = 1, color = color, iprint = 0, legend = [solvers[j]], show= False, xlabel='time', goal='minimum', name='nlp3')
    if solver == 'algencan':
        p.gradtol = 1e-1
    elif solver == 'ralg':
        p.debug = 1

    r = p.solve(solver, debug=1)

    results[solver] = (r.ff, p.getMaxResidual(r.xf), r.elapsed['solver_time'], r.elapsed['solver_cputime'], r.evals['f'], r.evals['c'], r.evals['h'])
    subplot(2,1,1)
    F0 = asscalar(p.f(p.x0))
    lines.append(plot([0, 1e-15], [F0, F0], color= colors[j]))

for i in range(2):
    subplot(2,1,i+1)
    legend(lines, solvers)

subplots_adjust(bottom=0.2, hspace=0.3)

xl = ['Solver                              f_opt     MaxConstr   Time   CPUTime  fEvals  cEvals  hEvals']
for i in range(len(results)):
    xl.append((expandtabs(ljust(solvers[i], 16)+' \t', 15)+'%0.2f'% (results[solvers[i]][0]) + '        %0.1e' % (results[solvers[i]][1]) + ('      %0.2f'% (results[solvers[i]][2])) + '    %0.2f      '% (results[solvers[i]][3]) + str(results[solvers[i]][4]) + '   ' + rjust(str(results[solvers[i]][5]), 5) + expandtabs('\t' +str(results[solvers[i]][6]),8)))

xl = '\n'.join(xl)
subplot(2,1,1)
xlabel('Time elapsed (without graphic output), sec')

from pylab import *
subplot(2,1,2)
xlabel(xl)

show()

